# app.py
from flask import Flask, request, jsonify
import os
import tempfile
import subprocess
import logging
from werkzeug.utils import secure_filename
import socket
import argparse
from flask import current_app

app = Flask(__name__)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 配置临时文件目录
TEMP_DIR = tempfile.gettempdir()
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}

import sys
import traceback
# 全局异常捕获
def global_except_hook(exctype, value, tb):
    print("⚠️ 未捕获的异常:", file=sys.stderr)
    traceback.print_exception(exctype, value, tb)
    sys.exit(1)

sys.excepthook = global_except_hook

def allowed_file(filename):
    """检查文件扩展名是否允许"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def find_free_port(start_port=5000):
    """查找可用端口"""
    port = start_port
    while port < 65535:
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.bind(('', port))
                return port
        except OSError:
            port += 1
    raise RuntimeError("No available ports found")

def run_inference(file_path, params):
    """运行推理脚本并返回结果"""
    try:
        # 从配置中获取脚本路径
        script_path = current_app.config['INFERENCE_SCRIPT']
        cmd = ["python", script_path, "--input", file_path]
        
        # 添加其他参数
        for key, value in params.items():
            if value is not None:  # 只添加非None的参数
                cmd.extend([f"--{key}", str(value)])
        
        # 运行推理脚本
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True
        )
        
        # 检查是否有错误输出
        if result.stderr:
            logger.warning(f"Script warning: {result.stderr}")
        
        # 尝试解析输出为JSON
        try:
            import json
            return json.loads(result.stdout)
        except json.JSONDecodeError:
            return {"output": result.stdout.strip()}
            
    except subprocess.CalledProcessError as e:
        logger.error(f"Script error: {e.stderr}")
        raise RuntimeError(f"Inference script failed: {e.stderr}")
    except Exception as e:
        logger.error(f"Error running inference: {str(e)}")
        raise

@app.route('/api/analyze', methods=['POST'])
def analyze_image():
    """处理图像分析请求"""
    try:
        # 检查是否有文件
        if 'file' not in request.files:
            return jsonify({'error': 'No file provided'}), 400
        
        file = request.files['file']
        if file.filename == '':
            return jsonify({'error': 'No file selected'}), 400
            
        if not allowed_file(file.filename):
            return jsonify({'error': 'File type not allowed'}), 400

        # 获取其他参数
        params = request.form.to_dict()
        
        # 安全地保存文件到临时目录
        filename = secure_filename(file.filename)
        temp_path = os.path.join(TEMP_DIR, filename)
        
        try:
            file.save(temp_path)
            logger.info(f"File saved temporarily at: {temp_path}")
            
            # 运行推理
            result = run_inference(temp_path, params)
            
            return jsonify({
                'status': 'success',
                'result': result
            })
            
        finally:
            # 清理临时文件
            try:
                if os.path.exists(temp_path):
                    os.remove(temp_path)
                    logger.info(f"Temporary file removed: {temp_path}")
            except Exception as e:
                logger.warning(f"Failed to remove temporary file: {str(e)}")
                
    except Exception as e:
        logger.error(f"Error processing request: {str(e)}")
        return jsonify({
            'status': 'error',
            'error': str(e)
        }), 500

def main():
    parser = argparse.ArgumentParser(description='Start the Benetech API server')
    parser.add_argument('--port', type=int, help='Port to run the server on')
    parser.add_argument('--host', type=str, default='0.0.0.0', help='Host to run the server on')
    parser.add_argument('--script_path', type=str, default='benetech_inference.py', help='Path to the inference script')
    args = parser.parse_args()
    # 将脚本路径存入 Flask 配置
    app.config['INFERENCE_SCRIPT'] = os.path.abspath(args.script_path)  # 转换为绝对路径

    # 如果没有指定端口，查找可用端口
    port = args.port if args.port else find_free_port()
    
    logger.info(f"Starting server on port {port}")
    app.run(host=args.host, port=port)

if __name__ == '__main__':
    main()
